{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ijGzTHJJUCPY"
      },
      "outputs": [],
      "source": [
        "# Copyright 2025 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JipjpNFHZ1NZ"
      },
      "source": [
        "# Imagen 3 Customized Avatar Images\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/vision/use-cases/imagen3_customization_avatar.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fvision%2Fuse-cases%2Fimagen3_customization_avatar.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Run in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/vision/use-cases/imagen3_customization_avatar.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>    \n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/vision/use-cases/imagen3_customization_avatar.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/vision/use-cases/imagen3_customization_avatar.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/vision/use-cases/imagen3_customization_avatar.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/vision/use-cases/imagen3_customization_avatar.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/vision/use-cases/imagen3_customization_avatar.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/vision/use-cases/imagen3_customization_avatar.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G1KDmM_PBAXz"
      },
      "source": [
        "| Author |\n",
        "| --- |\n",
        "| [Katie Nguyen](https://github.com/katiemn) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2c69e7975c15"
      },
      "source": [
        "## Overview\n",
        "\n",
        "### Imagen 3\n",
        "\n",
        "Imagen 3 on Vertex AI brings Google's generative AI image customization capabilities to application developers. It's capable of modifying images to fit a certain style. Thus, developers have more control when building next-generation AI products that transform their imagination into high quality visual assets. Learn more about [Imagen on Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/image/overview).\n",
        "\n",
        "In this tutorial, you will learn how to use the Google Gen AI SDK for Python to generate customized avatar images using few-shot learning with Imagen 3. You'll complete image preprocessing steps, supply a text prompt and guide new image generation in the following styles:\n",
        "\n",
        "- Watercolor\n",
        "- Pencil sketch\n",
        "- Marker illustration\n",
        "- 3D cartoon"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r11Gu7qNgx1p"
      },
      "source": [
        "## Get started\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Google Gen AI SDK for Python & OpenCV\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet google-genai opencv-python"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you are running this notebook on Google Colab, run the following cell to authenticate your environment.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ua6PDqB1iBSb"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "yTiDo0lRh6sc"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "import io\n",
        "\n",
        "from IPython.display import Markdown, display\n",
        "from PIL import Image as PIL_Image\n",
        "import cv2\n",
        "from google import genai\n",
        "import google.auth\n",
        "import google.auth.transport.requests\n",
        "from google.genai.types import (\n",
        "    ControlReferenceConfig,\n",
        "    ControlReferenceImage,\n",
        "    EditImageConfig,\n",
        "    Image,\n",
        "    Part,\n",
        "    SubjectReferenceConfig,\n",
        "    SubjectReferenceImage,\n",
        ")\n",
        "import matplotlib.pyplot as plt\n",
        "import requests"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information and create client\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
        "\n",
        "client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b84bdfdd7ed3"
      },
      "source": [
        "### Define helper functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "f936014f357d"
      },
      "outputs": [],
      "source": [
        "def send_request_to_google_api(api_endpoint, data=None):\n",
        "    \"\"\"\n",
        "    Sends an HTTP request to a Google API endpoint.\n",
        "\n",
        "    Args:\n",
        "        api_endpoint: The URL of the Google API endpoint.\n",
        "        data: (Optional) Dictionary of data to send in the request body (for POST, PUT, etc.).\n",
        "\n",
        "    Returns:\n",
        "        The response from the Google API.\n",
        "    \"\"\"\n",
        "\n",
        "    # Get access token calling API\n",
        "    creds, project = google.auth.default()\n",
        "    auth_req = google.auth.transport.requests.Request()\n",
        "    creds.refresh(auth_req)\n",
        "    access_token = creds.token\n",
        "\n",
        "    headers = {\n",
        "        \"Authorization\": f\"Bearer {access_token}\",\n",
        "        \"Content-Type\": \"application/json\",\n",
        "    }\n",
        "\n",
        "    response = requests.post(api_endpoint, headers=headers, json=data)\n",
        "    response.raise_for_status()\n",
        "    return response.json()\n",
        "\n",
        "\n",
        "def get_resized_bytes(image_bytes, max_size):\n",
        "    image = PIL_Image.open(io.BytesIO(image_bytes))\n",
        "    buffered = io.BytesIO()\n",
        "    image.thumbnail((max_size, max_size))\n",
        "    image.save(buffered, \"PNG\")\n",
        "    image.save(\"resized.png\", \"PNG\")\n",
        "    return buffered.getvalue()\n",
        "\n",
        "\n",
        "def segment_image(request: dict):\n",
        "    params = {\"mode\": \"semantic\", \"confidenceThreshold\": 0.5}\n",
        "    req = {\"instances\": [request], \"parameters\": params}\n",
        "    resp = send_request_to_google_api(segmentation_model, req)\n",
        "    return resp[\"predictions\"]\n",
        "\n",
        "\n",
        "def crop_and_pad_face(image_path, output_path, bbox, padding_factor=0.5):\n",
        "    image = cv2.imread(image_path)\n",
        "    H, W, _ = image.shape\n",
        "    y, x, h, w = bbox\n",
        "\n",
        "    pad_x = int(w * padding_factor)\n",
        "    pad_y = int(h * padding_factor)\n",
        "\n",
        "    crop_x1 = max(0, x - pad_x)\n",
        "    crop_y1 = max(0, y - pad_y)\n",
        "    crop_x2 = min(W, x + w + pad_x)\n",
        "    crop_y2 = min(H, y + h + pad_y)\n",
        "\n",
        "    cropped_image = image[crop_y1:crop_y2, crop_x1:crop_x2]\n",
        "    cv2.imwrite(output_path, cropped_image)\n",
        "\n",
        "\n",
        "def display_images(generated_image, ref_image) -> None:\n",
        "    fig, axis = plt.subplots(1, 2, figsize=(12, 6))\n",
        "    axis[0].imshow(generated_image)\n",
        "    axis[0].set_title(\"Modified Image\")\n",
        "    axis[1].imshow(ref_image)\n",
        "    axis[1].set_title(\"Reference Image\")\n",
        "    for ax in axis:\n",
        "        ax.axis(\"off\")\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VLmwIj2RD0Fx"
      },
      "source": [
        "### Load the  models\n",
        "\n",
        "Imagen 3 Customization: `imagen-3.0-capability-001`\n",
        "\n",
        "Image Segmentation: `image-segmentation-001`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "F-gd2ypQhh7K"
      },
      "outputs": [],
      "source": [
        "customization_model = \"imagen-3.0-capability-001\"\n",
        "segmentation_model = f\"https://us-central1-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/us-central1/publishers/google/models/image-segmentation-001:predict\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0ygjOGF3hJcz"
      },
      "source": [
        "## Image preprocessing\n",
        "\n",
        "In order to get the best quality out of your customized images, segmenting, facial detection, and cropping will help improve overall results."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iWoMrvxuhQS5"
      },
      "source": [
        "### Download an image\n",
        "\n",
        "For this process you'll need to use an image with one person that doesn't have an occluded face. You can download one to use locally from a Google Cloud Storage bucket or locally upload one to session storage.\n",
        "\n",
        "**Google Cloud Storage:** Switch out the URL in the `wget` command below to point to the URL of your image in a Google Cloud Storage bucket.\n",
        "\n",
        "**Session storage:** Click on the file icon on the left hand side. Then, click on the 'Upload to session storage' button and select an image from your computer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d3Bd9h2M1fhW"
      },
      "outputs": [],
      "source": [
        "!wget https://storage.googleapis.com/cloud-samples-data/generative-ai/image/man-in-field.png"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6iXy2ukoiduR"
      },
      "source": [
        "Enter the file name of the image you uploaded from Google Cloud Storage or your computer in the `IMAGE_FILE` variable below. By running the next cell you'll then rename your image file in Colab to be \"face.png\"."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "x1ZfjYh5P74D"
      },
      "outputs": [],
      "source": [
        "IMAGE_FILE = \"man-in-field.png\"  # @param {type: \"string\"}\n",
        "\n",
        "os.rename(IMAGE_FILE, \"face.png\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lVEWGTdDhVrc"
      },
      "source": [
        "### Image segmentation & background replacement\n",
        "\n",
        "\n",
        "> ⚠️ **Warning:** This process only works if the initial image has a single face that is not heavily occluded.\n",
        "\n",
        "In this next cell, you'll send a request to the image segmentation model to create a mask around the person in the image so that they can be placed on a solid white background."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1x1blbPnmXYY"
      },
      "outputs": [],
      "source": [
        "prompt = \"person\"\n",
        "with open(\"face.png\", \"rb\") as image_file:\n",
        "    image_bytes = image_file.read()\n",
        "\n",
        "new_bytes = get_resized_bytes(image_bytes, 640)\n",
        "base64_encoded_bytes = base64.b64encode(new_bytes)\n",
        "\n",
        "request = {\"prompt\": prompt, \"image\": {\"bytesBase64Encoded\": base64_encoded_bytes}}\n",
        "mask = segment_image(request)\n",
        "for m in mask:\n",
        "    bytes_b64 = dict(m)[\"bytesBase64Encoded\"]\n",
        "    decoded_image_data = base64.b64decode(bytes_b64)\n",
        "    image_stream = io.BytesIO(decoded_image_data)\n",
        "    PIL_Image.open(image_stream).save(\"mask.png\")\n",
        "\n",
        "display_images(PIL_Image.open(\"mask.png\"), PIL_Image.open(\"face.png\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eJa0q4numkjb"
      },
      "source": [
        "Now that you've segmented the person in the image, remove the background and replace it with a solid white backdrop."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A-PdxtELxIMb"
      },
      "outputs": [],
      "source": [
        "base_image = PIL_Image.open(\"resized.png\")\n",
        "pil_mask = PIL_Image.open(\"mask.png\")\n",
        "color = (255, 255, 255, 255)\n",
        "\n",
        "white_background = PIL_Image.new(\"RGBA\", base_image.size, color)\n",
        "white_background.paste(base_image, mask=pil_mask)\n",
        "\n",
        "white_background.save(\"remove-background.png\")\n",
        "display(white_background)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E3bUrTIyXIXE"
      },
      "source": [
        "### Facial detection\n",
        "\n",
        "Run the following code to detect the face in the image using the pre-trained Haar Cascade classifier that is built into OpenCV. This cell will also draw a green bounding box around the detected face.\n",
        "\n",
        "> ⚠️ **Warning:** If a face isn't detected, download a new photo and restart the image preprocessing steps. Otherwise, the rest of the tutorial will not work as expected."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DaBbxsxLfLpR"
      },
      "outputs": [],
      "source": [
        "img = cv2.imread(\"remove-background.png\")\n",
        "gray_image = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n",
        "\n",
        "# load the pre-trained Haar Cascade classifier that is built into OpenCV\n",
        "face_classifier = cv2.CascadeClassifier(\n",
        "    cv2.data.haarcascades + \"haarcascade_frontalface_default.xml\"\n",
        ")\n",
        "face = face_classifier.detectMultiScale(\n",
        "    gray_image, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30)\n",
        ")\n",
        "\n",
        "for x, y, w, h in face:\n",
        "    detected_face = [y, x, h, w]\n",
        "    box_2d_list = [int(x) for x in detected_face]\n",
        "\n",
        "for x, y, w, h in face:\n",
        "    cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 4)\n",
        "img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
        "\n",
        "plt.figure(figsize=(20, 10))\n",
        "plt.imshow(img_rgb)\n",
        "plt.axis(\"off\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cHIcGinHhb9h"
      },
      "source": [
        "### Crop the image\n",
        "\n",
        "Once a face is detected, crop around the face using a helper function defined in the beginning of this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "L1VAx5Vx_oia"
      },
      "outputs": [],
      "source": [
        "crop_and_pad_face(\"remove-background.png\", \"cropped_image.png\", box_2d_list)\n",
        "display(PIL_Image.open(\"cropped_image.png\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aKiU4Rkfhqks"
      },
      "source": [
        "## Create avatar images\n",
        "\n",
        "Now that you've preprocessed the initial image, it's time to customize the detected face in certain styles. Before you do this, you'll need to use Gemini to write a description of the person that can be used in later calls to Imagen."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FT1h_ttD2JU_"
      },
      "source": [
        "### Use Gemini to generate a description"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5EChEA82370I"
      },
      "outputs": [],
      "source": [
        "with open(\"cropped_image.png\", \"rb\") as f:\n",
        "    image = f.read()\n",
        "\n",
        "response = client.models.generate_content(\n",
        "    model=\"gemini-2.5-flash\",\n",
        "    contents=[\n",
        "        Part.from_bytes(data=image, mime_type=\"image/png\"),\n",
        "        \"Briefly provide a description of the person's face in this image with 6 words max, no punctuation, all lowercase\",\n",
        "    ],\n",
        ")\n",
        "\n",
        "display(Markdown(response.text))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "odj4xYgPhzUC"
      },
      "source": [
        "### Style 1: Watercolor"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8pyAJlvQsocc"
      },
      "outputs": [],
      "source": [
        "subject_image = Image.from_file(location=\"cropped_image.png\")\n",
        "\n",
        "subject_reference_image = SubjectReferenceImage(\n",
        "    reference_id=1,\n",
        "    reference_image=subject_image,\n",
        "    config=SubjectReferenceConfig(\n",
        "        subject_description=response.text, subject_type=\"SUBJECT_TYPE_PERSON\"\n",
        "    ),\n",
        ")\n",
        "\n",
        "control_reference_image = ControlReferenceImage(\n",
        "    reference_id=2,\n",
        "    reference_image=subject_image,\n",
        "    config=ControlReferenceConfig(control_type=\"CONTROL_TYPE_FACE_MESH\"),\n",
        ")\n",
        "\n",
        "prompt = f\"Create a watercolor image of a {response.text} [1] in the pose of the control image [2] to match the description: A watercolor portrait of {response.text} [1] in a watercolor style, light and low-contrast color stokes, bright pastel color, a warm atmosphere, clean background with watercolor brushstrokes, bold dry brush, dry on dry technique, grainy textured paper, contrasting light and shadow, fine watery brush strokes, patchy details\"\n",
        "\n",
        "image = client.models.edit_image(\n",
        "    model=customization_model,\n",
        "    prompt=prompt,\n",
        "    reference_images=[subject_reference_image, control_reference_image],\n",
        "    config=EditImageConfig(\n",
        "        edit_mode=\"EDIT_MODE_DEFAULT\",\n",
        "        number_of_images=1,\n",
        "        safety_filter_level=\"BLOCK_MEDIUM_AND_ABOVE\",\n",
        "        person_generation=\"ALLOW_ADULT\",\n",
        "    ),\n",
        ")\n",
        "\n",
        "display_images(\n",
        "    image.generated_images[0].image._pil_image, PIL_Image.open(\"resized.png\")\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VWeJlcrC4AH0"
      },
      "source": [
        "### Style 2: Pencil sketch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TealmVHdCfrE"
      },
      "outputs": [],
      "source": [
        "subject_image = Image.from_file(location=\"cropped_image.png\")\n",
        "\n",
        "subject_reference_image = SubjectReferenceImage(\n",
        "    reference_id=1,\n",
        "    reference_image=subject_image,\n",
        "    config=SubjectReferenceConfig(\n",
        "        subject_description=response.text, subject_type=\"SUBJECT_TYPE_PERSON\"\n",
        "    ),\n",
        ")\n",
        "\n",
        "control_reference_image = ControlReferenceImage(\n",
        "    reference_id=2,\n",
        "    reference_image=subject_image,\n",
        "    config=ControlReferenceConfig(control_type=\"CONTROL_TYPE_FACE_MESH\"),\n",
        ")\n",
        "\n",
        "prompt = f\"Create a pencil sketch of {response.text} [1] in the pose of the control image [2] to match the description: A simple pencil style sketch of a portrait of {response.text} with 6B and graphite pencils, white background, high quality, visible pencil lines, looking at the camera, natural human eyes\"\n",
        "\n",
        "image = client.models.edit_image(\n",
        "    model=customization_model,\n",
        "    prompt=prompt,\n",
        "    reference_images=[subject_reference_image, control_reference_image],\n",
        "    config=EditImageConfig(\n",
        "        edit_mode=\"EDIT_MODE_DEFAULT\",\n",
        "        number_of_images=1,\n",
        "        safety_filter_level=\"BLOCK_MEDIUM_AND_ABOVE\",\n",
        "        person_generation=\"ALLOW_ADULT\",\n",
        "    ),\n",
        ")\n",
        "\n",
        "display_images(\n",
        "    image.generated_images[0].image._pil_image, PIL_Image.open(\"resized.png\")\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "84i1WBCd4GAy"
      },
      "source": [
        "### Style 3: Illustration"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HLYGA-VfDEk0"
      },
      "outputs": [],
      "source": [
        "subject_image = Image.from_file(location=\"cropped_image.png\")\n",
        "\n",
        "subject_reference_image = SubjectReferenceImage(\n",
        "    reference_id=1,\n",
        "    reference_image=subject_image,\n",
        "    config=SubjectReferenceConfig(\n",
        "        subject_description=response.text, subject_type=\"SUBJECT_TYPE_PERSON\"\n",
        "    ),\n",
        ")\n",
        "\n",
        "control_reference_image = ControlReferenceImage(\n",
        "    reference_id=2,\n",
        "    reference_image=subject_image,\n",
        "    config=ControlReferenceConfig(control_type=\"CONTROL_TYPE_FACE_MESH\"),\n",
        ")\n",
        "\n",
        "prompt = f\"Create an illustration of {response.text} [1] in the pose of the control image [2] to match the description: a portrait of {response.text} [1] in a vector illustration style with bold outlines, pastel tones, clean and precise lines, organic wide brushstrokes in the background, soft aesthetics theme, blended color on the skin, dramatic lighting and contrasting highlights and shadows\"\n",
        "\n",
        "image = client.models.edit_image(\n",
        "    model=customization_model,\n",
        "    prompt=prompt,\n",
        "    reference_images=[subject_reference_image, control_reference_image],\n",
        "    config=EditImageConfig(\n",
        "        edit_mode=\"EDIT_MODE_DEFAULT\",\n",
        "        number_of_images=1,\n",
        "        safety_filter_level=\"BLOCK_MEDIUM_AND_ABOVE\",\n",
        "        person_generation=\"ALLOW_ADULT\",\n",
        "    ),\n",
        ")\n",
        "\n",
        "display_images(\n",
        "    image.generated_images[0].image._pil_image, PIL_Image.open(\"resized.png\")\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Moa6b7IA4KJh"
      },
      "source": [
        "### Style 4: 3D cartoon"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hnx3sen8DYxc"
      },
      "outputs": [],
      "source": [
        "subject_image = Image.from_file(location=\"cropped_image.png\")\n",
        "\n",
        "subject_reference_image = SubjectReferenceImage(\n",
        "    reference_id=1,\n",
        "    reference_image=subject_image,\n",
        "    config=SubjectReferenceConfig(\n",
        "        subject_description=response.text, subject_type=\"SUBJECT_TYPE_PERSON\"\n",
        "    ),\n",
        ")\n",
        "\n",
        "control_reference_image = ControlReferenceImage(\n",
        "    reference_id=2,\n",
        "    reference_image=subject_image,\n",
        "    config=ControlReferenceConfig(control_type=\"CONTROL_TYPE_FACE_MESH\"),\n",
        ")\n",
        "\n",
        "prompt = f\"Create a 3D cartoon style image of {response.text} [1] in the pose of the control image [2] to match the description: a portrait of  {response.text} [1] in 3D cartoon style with a blurred background. A cute and lovely character, smiley face looking at the camera, pastel color tones, high quality, super details, skin texture, texture mapping, soft shadows, soft realistic lighting, vibrant colors\"\n",
        "\n",
        "image = client.models.edit_image(\n",
        "    model=customization_model,\n",
        "    prompt=prompt,\n",
        "    reference_images=[subject_reference_image, control_reference_image],\n",
        "    config=EditImageConfig(\n",
        "        edit_mode=\"EDIT_MODE_DEFAULT\",\n",
        "        number_of_images=1,\n",
        "        safety_filter_level=\"BLOCK_MEDIUM_AND_ABOVE\",\n",
        "        person_generation=\"ALLOW_ADULT\",\n",
        "    ),\n",
        ")\n",
        "\n",
        "display_images(\n",
        "    image.generated_images[0].image._pil_image, PIL_Image.open(\"resized.png\")\n",
        ")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "imagen3_customization_avatar.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
